#pragma once

#include <NvInfer.h>
#include <NvInferRuntimeCommon.h>

#include "logger.h"
#include "opencv2/opencv.hpp"
#include "spdlog/spdlog.h"
#include "yololayer.h"

using namespace nvinfer1;

namespace TRT {

class TRTDeleter {
 public:
  template <typename T>
  void operator()(T *obj) const {
    if (obj) {
#if NV_TENSORRT_SONAME_MAJOR == 7
      obj->destroy();
#elif NV_TENSORRT_SONAME_MAJOR == 8
      delete obj;
#endif
      SPDLOG_DEBUG("[TRTDeleter] destroy.");
    }
  }
};

static auto StreamDeleter = [](cudaStream_t *pStream) {
  if (pStream) {
    cudaStreamDestroy(*pStream);
    delete pStream;
  }
};

#if NV_TENSORRT_SONAME_MAJOR == 7
class TRTLogger : public nvinfer1::ILogger {
#elif NV_TENSORRT_SONAME_MAJOR == 8
class TRTLogger : public sample::Logger {
#endif

 public:
  TRTLogger() = default;
  ~TRTLogger() = default;

  void log(Severity severity, const char *msg) noexcept;
  int GetVerbosity();
};

}  // namespace TRT
namespace component {

enum class Model {
  NoMask = 0,
  WithMask = 1,
};

std::string ModelToString(Model model);

void Draw(const cv::Mat &output, TRT::Detection detection, int id = 0);

void PrintInfo(TRT::Detection detection);

std::string Replace(const std::string orin, const std::string &add);

}  // namespace component

template <typename T>
using UniquePtr = std::unique_ptr<T, TRT::TRTDeleter>;